package com.gitee.wsl.ext.coroutines

import kotlinx.atomicfu.locks.ReentrantLock
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.coroutines.CoroutineContext
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.currentCoroutineContext
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract


fun Mutex.lockBlocking(owner: Any? = null) {
    while (!tryLock(owner));
}

inline fun <T> Mutex.withLockBlocking(owner: Any? = null, action: () -> T): T {
    try {
        lockBlocking(owner)
        return action()
    } finally {
        unlock(owner)
    }
}

inline operator fun <T> ReentrantLock.invoke(callback: () -> T): T {
    lock()
    try {
        return callback()
    } finally {
        unlock()
    }
}

// Based on https://gist.github.com/elizarov/9a48b9709ffd508909d34fab6786acfe
// See also https://elizarov.medium.com/phantom-of-the-coroutine-afc63b03a131
suspend fun <T> Mutex.withReentrantLock(block: suspend () -> T): T {
    val key = ReentrantMutexContextKey(this)
    // call block directly when this mutex is already locked in the context
    if (currentCoroutineContext()[key] != null) return block()
    // otherwise add it to the context and lock the mutex
    return withContext(ReentrantMutexContextElement(key)) {
        withLock { block() }
    }
}

internal class ReentrantMutexContextElement(
    override val key: ReentrantMutexContextKey,
) : CoroutineContext.Element

internal data class ReentrantMutexContextKey(
    val mutex: Mutex,
) : CoroutineContext.Key<ReentrantMutexContextElement>


/**
 * Executes the given [action] within a [CoroutineScope] under this [Mutex]'s lock.
 *
 * @param owner Optional owner token for debugging.
 * When [owner] is specified (non-`null` value) and this mutex is already locked with the same (as in `===`) token,
 * this function throws an [IllegalStateException].
 *
 * @since 0.0.1-alpha.4
 * @see coroutineScope
 * @see Mutex.withLock
 */
@Throws(IllegalStateException::class, CancellationException::class)
@OptIn(ExperimentalContracts::class)
public suspend fun <R> Mutex.coroutineScopeWithLock(owner: Any? = null, action: suspend CoroutineScope.() -> R): R {
    contract {
        callsInPlace(lambda = action, kind = InvocationKind.EXACTLY_ONCE)
    }

    return withLock(owner) { coroutineScope(block = action) }
}

/**
 * Executes the given [action], and if this [Mutex] is not `null`, under its lock.
 *
 * @param owner Optional owner token for debugging.
 * When [owner] is specified (non-`null` value) and this mutex is already locked with the same (as in `===`) token,
 * this function throws an [IllegalStateException].
 *
 * @since 0.0.1-alpha.4
 * @see Mutex.withLock
 */
@OptIn(ExperimentalContracts::class)
public suspend inline fun <R> Mutex?.withLockIfAny(owner: Any? = null, action: () -> R): R {
    contract {
        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
    }

    return if (this == null) action() else withLock(owner, action)
}

/**
 * Executes the given [action] within a [CoroutineScope], and if this [Mutex] is not `null`, under its lock.
 *
 * @param owner Optional owner token for debugging.
 * When [owner] is specified (non-`null` value) and this mutex is already locked with the same (as in `===`) token,
 * this function throws an [IllegalStateException].
 *
 * @since 0.0.1-alpha.4
 * @see coroutineScope
 * @see Mutex.withLock
 * @see Mutex.withLockIfAny
 */
@OptIn(ExperimentalContracts::class)
public suspend fun <R> Mutex?.coroutineScopeWithLockIfAny(owner: Any? = null, action: suspend CoroutineScope.() -> R): R {
    contract {
        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
    }

    return withLockIfAny(owner) { coroutineScope(block = action) }
}
